/*
 * Copyright (C) 2008 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.dx.ssa;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

import com.android.dx.rop.code.LocalItem;
import com.android.dx.rop.code.PlainCstInsn;
import com.android.dx.rop.code.PlainInsn;
import com.android.dx.rop.code.RegOps;
import com.android.dx.rop.code.RegisterSpec;
import com.android.dx.rop.code.RegisterSpecList;
import com.android.dx.rop.code.Rop;
import com.android.dx.rop.code.Rops;
import com.android.dx.rop.code.SourcePosition;
import com.android.dx.rop.code.ThrowingCstInsn;
import com.android.dx.rop.cst.Constant;
import com.android.dx.rop.cst.CstString;
import com.android.dx.rop.cst.TypedConstant;
import com.android.dx.rop.type.StdTypeList;
import com.android.dx.rop.type.TypeBearer;

/**
 * Collects constants that are used more than once at the top of the method
 * block. This increases register usage by about 5% but decreases insn size by
 * about 3%.
 */
public class ConstCollector {

	/** Maximum constants to collect per method. Puts cap on reg use */
	private static final int MAX_COLLECTED_CONSTANTS = 5;

	/**
	 * Also collect string consts, although they can throw exceptions. This is
	 * off now because it just doesn't seem to gain a whole lot. TODO if you
	 * turn this on, you must change SsaInsn.hasSideEffect() to return false for
	 * const-string insns whose exceptions are not caught in the current method.
	 */
	private static boolean COLLECT_STRINGS = false;

	/**
	 * If true, allow one local var to be involved with a collected const.
	 * Turned off because it mostly just inserts more moves.
	 */
	private static boolean COLLECT_ONE_LOCAL = false;

	/** method we're processing */
	private final SsaMethod ssaMeth;

	/**
	 * Processes a method.
	 * 
	 * @param ssaMethod
	 *            {@code non-null;} method to process
	 */
	public static void process(SsaMethod ssaMethod) {
		ConstCollector cc = new ConstCollector(ssaMethod);
		cc.run();
	}

	/**
	 * Constructs an instance.
	 * 
	 * @param ssaMethod
	 *            {@code non-null;} method to process
	 */
	private ConstCollector(SsaMethod ssaMethod) {
		this.ssaMeth = ssaMethod;
	}

	/**
	 * Applies the optimization.
	 */
	private void run() {
		int regSz = ssaMeth.getRegCount();

		ArrayList<TypedConstant> constantList = getConstsSortedByCountUse();

		int toCollect = Math.min(constantList.size(), MAX_COLLECTED_CONSTANTS);

		SsaBasicBlock start = ssaMeth.getEntryBlock();

		// Constant to new register containing the constant
		HashMap<TypedConstant, RegisterSpec> newRegs = new HashMap<TypedConstant, RegisterSpec>(
				toCollect);

		for (int i = 0; i < toCollect; i++) {
			TypedConstant cst = constantList.get(i);
			RegisterSpec result = RegisterSpec.make(ssaMeth.makeNewSsaReg(),
					cst);

			Rop constRop = Rops.opConst(cst);

			if (constRop.getBranchingness() == Rop.BRANCH_NONE) {
				start.addInsnToHead(new PlainCstInsn(Rops.opConst(cst),
						SourcePosition.NO_INFO, result, RegisterSpecList.EMPTY,
						cst));
			} else {
				// We need two new basic blocks along with the new insn
				SsaBasicBlock entryBlock = ssaMeth.getEntryBlock();
				SsaBasicBlock successorBlock = entryBlock.getPrimarySuccessor();

				// Insert a block containing the const insn.
				SsaBasicBlock constBlock = entryBlock
						.insertNewSuccessor(successorBlock);

				constBlock.replaceLastInsn(new ThrowingCstInsn(constRop,
						SourcePosition.NO_INFO, RegisterSpecList.EMPTY,
						StdTypeList.EMPTY, cst));

				// Insert a block containing the move-result-pseudo insn.

				SsaBasicBlock resultBlock = constBlock
						.insertNewSuccessor(successorBlock);
				PlainInsn insn = new PlainInsn(Rops.opMoveResultPseudo(result
						.getTypeBearer()), SourcePosition.NO_INFO, result,
						RegisterSpecList.EMPTY);

				resultBlock.addInsnToHead(insn);
			}

			newRegs.put(cst, result);
		}

		updateConstUses(newRegs, regSz);
	}

	/**
	 * Gets all of the collectable constant values used in this method, sorted
	 * by most used first. Skips non-collectable consts, such as non-string
	 * object constants
	 * 
	 * @return {@code non-null;} list of constants in most-to-least used order
	 */
	private ArrayList<TypedConstant> getConstsSortedByCountUse() {
		int regSz = ssaMeth.getRegCount();

		final HashMap<TypedConstant, Integer> countUses = new HashMap<TypedConstant, Integer>();

		/*
		 * Each collected constant can be used by just one local (used only if
		 * COLLECT_ONE_LOCAL is true).
		 */
		final HashSet<TypedConstant> usedByLocal = new HashSet<TypedConstant>();

		// Count how many times each const value is used.
		for (int i = 0; i < regSz; i++) {
			SsaInsn insn = ssaMeth.getDefinitionForRegister(i);

			if (insn == null || insn.getOpcode() == null)
				continue;

			RegisterSpec result = insn.getResult();
			TypeBearer typeBearer = result.getTypeBearer();

			if (!typeBearer.isConstant())
				continue;

			TypedConstant cst = (TypedConstant) typeBearer;

			// Find defining instruction for move-result-pseudo instructions
			if (insn.getOpcode().getOpcode() == RegOps.MOVE_RESULT_PSEUDO) {
				int pred = insn.getBlock().getPredecessors().nextSetBit(0);
				ArrayList<SsaInsn> predInsns;
				predInsns = ssaMeth.getBlocks().get(pred).getInsns();
				insn = predInsns.get(predInsns.size() - 1);
			}

			if (insn.canThrow()) {
				/*
				 * Don't move anything other than strings -- the risk of
				 * changing where an exception is thrown is too high.
				 */
				if (!(cst instanceof CstString) || !COLLECT_STRINGS) {
					continue;
				}
				/*
				 * We can't move any throwable const whose throw will be caught,
				 * so don't count them.
				 */
				if (insn.getBlock().getSuccessors().cardinality() > 1) {
					continue;
				}
			}

			/*
			 * TODO: Might be nice to try and figure out which local wins most
			 * when collected.
			 */
			if (ssaMeth.isRegALocal(result)) {
				if (!COLLECT_ONE_LOCAL) {
					continue;
				} else {
					if (usedByLocal.contains(cst)) {
						// Count one local usage only.
						continue;
					} else {
						usedByLocal.add(cst);
					}
				}
			}

			Integer has = countUses.get(cst);
			if (has == null) {
				countUses.put(cst, 1);
			} else {
				countUses.put(cst, has + 1);
			}
		}

		// Collect constants that have been reused.
		ArrayList<TypedConstant> constantList = new ArrayList<TypedConstant>();
		for (Map.Entry<TypedConstant, Integer> entry : countUses.entrySet()) {
			if (entry.getValue() > 1) {
				constantList.add(entry.getKey());
			}
		}

		// Sort by use, with most used at the beginning of the list.
		Collections.sort(constantList, new Comparator<Constant>() {

			public int compare(Constant a, Constant b) {
				int ret;
				ret = countUses.get(b) - countUses.get(a);

				if (ret == 0) {
					/*
					 * Provide sorting determinisim for constants with same
					 * usage count.
					 */
					ret = a.compareTo(b);
				}

				return ret;
			}

			@Override
			public boolean equals(Object obj) {
				return obj == this;
			}
		});

		return constantList;
	}

	/**
	 * Inserts mark-locals if necessary when changing a register. If the
	 * definition of {@code origReg} is associated with a local variable, then
	 * insert a mark-local for {@code newReg} just below it. We expect the
	 * definition of {@code origReg} to ultimately be removed by the dead code
	 * eliminator
	 * 
	 * @param origReg
	 *            {@code non-null;} original register
	 * @param newReg
	 *            {@code non-null;} new register that will replace
	 *            {@code origReg}
	 */
	private void fixLocalAssignment(RegisterSpec origReg, RegisterSpec newReg) {
		for (SsaInsn use : ssaMeth.getUseListForRegister(origReg.getReg())) {
			RegisterSpec localAssignment = use.getLocalAssignment();
			if (localAssignment == null) {
				continue;
			}

			if (use.getResult() == null) {
				/*
				 * This is a mark-local. it will be updated when all uses are
				 * updated.
				 */
				continue;
			}

			LocalItem local = localAssignment.getLocalItem();

			// Un-associate original use.
			use.setResultLocal(null);

			// Now add a mark-local to the new reg immediately after.
			newReg = newReg.withLocalItem(local);

			SsaInsn newInsn = SsaInsn.makeFromRop(
					new PlainInsn(Rops.opMarkLocal(newReg),
							SourcePosition.NO_INFO, null, RegisterSpecList
									.make(newReg)), use.getBlock());

			ArrayList<SsaInsn> insns = use.getBlock().getInsns();

			insns.add(insns.indexOf(use) + 1, newInsn);
		}
	}

	/**
	 * Updates all uses of various consts to use the values in the newly
	 * assigned registers.
	 * 
	 * @param newRegs
	 *            {@code non-null;} mapping between constant and new reg
	 * @param origRegCount
	 *            {@code >=0;} original SSA reg count, not including newly added
	 *            constant regs
	 */
	private void updateConstUses(HashMap<TypedConstant, RegisterSpec> newRegs,
			int origRegCount) {

		/*
		 * set of constants associated with a local variable; used only if
		 * COLLECT_ONE_LOCAL is true.
		 */
		final HashSet<TypedConstant> usedByLocal = new HashSet<TypedConstant>();

		final ArrayList<SsaInsn>[] useList = ssaMeth.getUseListCopy();

		for (int i = 0; i < origRegCount; i++) {
			SsaInsn insn = ssaMeth.getDefinitionForRegister(i);

			if (insn == null) {
				continue;
			}

			final RegisterSpec origReg = insn.getResult();
			TypeBearer typeBearer = insn.getResult().getTypeBearer();

			if (!typeBearer.isConstant())
				continue;

			TypedConstant cst = (TypedConstant) typeBearer;
			final RegisterSpec newReg = newRegs.get(cst);

			if (newReg == null) {
				continue;
			}

			if (ssaMeth.isRegALocal(origReg)) {
				if (!COLLECT_ONE_LOCAL) {
					continue;
				} else {
					/*
					 * TODO: If the same local gets the same cst multiple times,
					 * it would be nice to reuse the register.
					 */
					if (usedByLocal.contains(cst)) {
						continue;
					} else {
						usedByLocal.add(cst);
						fixLocalAssignment(origReg, newRegs.get(cst));
					}
				}
			}

			// maps an original const register to the new collected register
			RegisterMapper mapper = new RegisterMapper() {

				@Override
				public int getNewRegisterCount() {
					return ssaMeth.getRegCount();
				}

				@Override
				public RegisterSpec map(RegisterSpec registerSpec) {
					if (registerSpec.getReg() == origReg.getReg()) {
						return newReg
								.withLocalItem(registerSpec.getLocalItem());
					}

					return registerSpec;
				}
			};

			for (SsaInsn use : useList[origReg.getReg()]) {
				if (use.canThrow()
						&& use.getBlock().getSuccessors().cardinality() > 1) {
					continue;
				}
				use.mapSourceRegisters(mapper);
			}
		}
	}
}
